
from grad_example import *
import sympy as sp
import re

import warnings
warnings.filterwarnings("ignore")

def mat_get(m, x):
    A = sp.zeros(m,m)
    for i in range(1, m+1):
        for j in range(1, m+1):
            pij = sp.Symbol('p_'+str(i)+str(j), positive=True)
            if i == j:
                A[i-1,j-1] = pij
            else:
                k = min(i, j)
                t = max(i, j)
                dij = sp.Symbol('d_'+str(k)+str(t), positive=True)
                A[i-1,j-1] = pij*sp.Pow(x, dij)
    return A



# def c1(m, N): 
#     x = sp.Symbol('x', positive=True)
#     a = sp.Wild('a')
#     b = sp.Wild('b')
    
#     s = sp.zeros(m, 1)
#     for i in range(1, m+1):
#         pij = sp.Symbol('p_a'+str(i), positive=True)
#         dij = sp.Symbol('d_a'+str(i), positive=True)
#         s[i-1,0] = pij*sp.Pow(x, dij)
    
#     s = s.T
#     A = mat_get(m, x)
#     for i in range(N-4):
#         A = A*A
    
#     t = sp.ones(m, 1)
#     A = s*A*t
#     func_symbol = A[0,0]
#     func_symbol = sp.expand(func_symbol)
#     func_symbol = sp.powsimp(func_symbol, force=True)
#     func_symbol = sp.collect(func_symbol, a**b)
    
#     # sm = re.findall('p[_ap0-9+ \*]+[0-9]', str(func_symbol))
#     # return len(sm)
#     return func_symbol

# # for N in range(4, 12):
# #     print(c1(2, N))

# def c2(m, N): 
#     x = sp.Symbol('x', positive=True)
#     a = sp.Wild('a')
#     b = sp.Wild('b')
    
#     s = sp.zeros(m, 1)
#     for i in range(1, m+1):
#         pij = sp.Symbol('p_a'+str(i), positive=True)
#         dij = sp.Symbol('d_a'+str(i), positive=True)
#         s[i-1,0] = pij*sp.Pow(x, dij)
    
#     s = s.T
#     A = mat_get(m, x)
#     for i in range(N-4):
#         A = A*A
    
#     t = sp.ones(m, 1)
#     for i in range(1, m+1):
#         dij = sp.Symbol('d_'+str(i)+'b', positive=True)
#         t[i-1,0] = sp.Pow(x, dij)

#     A = s*A*t
#     func_symbol = A[0,0]
#     func_symbol = sp.expand(func_symbol)
#     func_symbol = sp.powsimp(func_symbol, force=True)
#     func_symbol = sp.collect(func_symbol, a**b)
    
#     # sm = re.findall('p[_ap0-9+ \*]+[0-9]', str(func_symbol))
#     # return len(sm)
#     return func_symbol

m = 3
x = sp.Symbol('x', positive=True)
a = sp.Wild('a')
b = sp.Wild('b')

A = mat_get(m, x)
s1 = sp.zeros(m, 1)
for i in range(1, m+1):
    pij = sp.Symbol('p_a'+str(i), positive=True)
    dij = sp.Symbol('d_a'+str(i), positive=True)
    s1[i-1,0] = pij*sp.Pow(x, dij)
s1 = s1.T
s2 = sp.zeros(m, 1)
for i in range(1, m+1):
    dij = sp.Symbol('d_b'+str(i), positive=True)
    s2[i-1,0] = sp.Pow(x, dij)
s3 = sp.zeros(m, 1)
for i in range(1, m+1):
    pij = sp.Symbol('p_b'+str(i), positive=True)
    dij = sp.Symbol('d_b'+str(i), positive=True)
    s3[i-1,0] = pij*sp.Pow(x, dij)
s3 = s3.T
s4 = sp.ones(m, 1)
for i in range(1, m+1):
    dij = sp.Symbol('d_c'+str(i), positive=True)
    s4[i-1,0] = sp.Pow(x, dij)

s5 = sp.zeros(m, 1)
for i in range(1, m+1):
    pij = sp.Symbol('p_c'+str(i), positive=True)
    dij = sp.Symbol('d_c'+str(i), positive=True)
    s5[i-1,0] = pij*sp.Pow(x, dij)
s5 = s5.T
s6 = sp.ones(m, 1)
# for i in range(1, m+1):
#     dij = sp.Symbol('d_'+str(i)+'c', positive=True)
#     s6[i-1,0] = sp.Pow(x, dij)
# s7 = sp.zeros(m, 1)
# for i in range(1, m+1):
#     pij = sp.Symbol('p_0'+str(i), positive=True)
#     s7[i-1,0] = pij
# s7 = s7.T
# s8 = sp.zeros(m, 1)
# for i in range(1, m+1):
#     dij = sp.Symbol('d_a'+str(i), positive=True)
#     s8[i-1,0] = sp.Pow(x, dij)


A = s1*s2*s3*A*A*s4*s5*s6
func_symbol = A[0,0]
func_symbol = sp.expand(func_symbol)
func_symbol = sp.powsimp(func_symbol, force=True)
func_symbol = sp.collect(func_symbol, a**b)


def gf(func_symbol):
    s = re.findall('[p0-9][\*_][_eabcp0-9+ \*]+p_[aebc0-9][0-9]', str(func_symbol))
    H = 0
    for i in s:
        t = sp.sympify(i)
        H = H + t*sp.log(t)
    return H

def sl(H):
    symbol_set = H.free_symbols
    symbol_list = list(symbol_set)
    a = []
    for i in symbol_list:
        j = str(i)
        a.append(j)
    a.sort()
    symbol_list = []
    for i in a:
        j = sp.symbols(i)
        symbol_list.append(j)
    
    return symbol_list

# func_symbol = c1(4, 6)
H = gf(func_symbol)
symbol_list = sl(H)
H = -H/sp.log(2)

subs_dict = {}
remove_list = []
for i in range(2, len(symbol_list), 3):
    subs_dict[symbol_list[i]] = 1 - symbol_list[i-1] - symbol_list[i-2]
    remove_list.append(symbol_list[i])

H = H.subs(subs_dict)
for i in remove_list:
    symbol_list.remove(i)

# # # 变量映射转换
# # l = len(symbol_list)
# # x_list = list(sp.symbols('x:'+str(l)))
# # subs_dict = {}
# # for i in range(l):
# #     subs_dict[symbol_list[i]] = 1/(1+x_list[i]*x_list[i])
# # H = H.subs(subs_dict)

# # ans = grad_up(H, x_list, [3,3,3,3,3,3,3,3,3,3])
# # print(ans)

ans = grad_up(H, symbol_list, [0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1])
print(ans)